import time

from loguru import logger


class RequestTimeMiddleware:
    def __init__(self, get_response):
        """初始化中间件，接收 get_response 函数"""
        self.get_response = get_response
        self.start_time = None
        self.end_time = None

    def __call__(self, request):
        self.start_time = time.time()
        response = self.get_response(request)
        self.end_time = time.time()
        spend_time = (self.end_time - self.start_time) * 1000
        logger.info(f"请求耗时：{spend_time:.2f}ms")
        response["X-Request-Duration-MS"] = f"{spend_time:.2f}"
        return response
